package v1

import (
	"context"
	"fmt"
	"log/slog"
	llm "quipus/internal/llm"
	v1pb "quipus/proto/gen/api/v1"
	cpb "quipus/proto/gen/common"
	"strings"
	"time"

	"github.com/gofiber/websocket/v2"
	"github.com/lithammer/shortuuid/v4"
	"google.golang.org/protobuf/proto"
)

func (s *APIV1Service) handleCommonWebSocketMessages(ctx context.Context, cancel context.CancelFunc, tid string, conn *websocket.Conn, ragMode cpb.RagMode, inMessageChan chan *v1pb.Exchange) {
	writeMutex := s.WsManager.MutexMap[tid]

	sendInputMessage := func(askId string, content string, phase v1pb.Phase, mode cpb.RagMode) {
		inMsg := &v1pb.Exchange{
			AskId:   askId,
			Content: content,
			Phase:   phase,
			Mode:    mode,
		}
		select {
		case inMessageChan <- inMsg:
			slog.Info("Message sent to inMessageChan", "message", content)
		case <-ctx.Done():
			return
		default:
			slog.Warn("Failed to send message to inMessageChan: channel full or closed")
		}
	}
	for {
		select {
		case <-ctx.Done():
			return
		default:
			_, msg, err := conn.ReadMessage()
			if err != nil {
				if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
					slog.Warn("WebSocket close error:", slog.Any("err", err))
					cancel()
				} else {
					slog.Warn("WebSocket other error:", slog.Any("err", err))
					cancel()
				}
				return
			}

			if string(msg) == "xxpingxx" {
				sendpong := func() {
					writeMutex.Lock()
					defer writeMutex.Unlock()

					err := conn.WriteMessage(websocket.TextMessage, []byte("xxpongxx"))
					if err != nil {
						slog.Warn("Error sending pong:", slog.Any("err", err))
					}
				}
				sendpong()
				continue
			}

			var chatMessage v1pb.Exchange
			if err := proto.Unmarshal(msg, &chatMessage); err != nil {
				slog.Warn("Error unmarshaling message:", slog.Any("err", err))
				continue
			}
			ragMode := ragMode
			if chatMessage.Mode != cpb.RagMode_RAG_MODE_UNSPECIFIED {
				ragMode = chatMessage.Mode
			}
			if chatMessage.Phase == v1pb.Phase_HELLO {
				content := chatMessage.Content
				if content == "" {
					content = s.WsManager.InitMessage[chatMessage.AskId]
					if content == "" {
						slog.Info("hello phase message is empty ", chatMessage.AskId, content)
						continue
					}
					slog.Info("hello init granary message", chatMessage.AskId, content)
					sendInputMessage(chatMessage.AskId, content, chatMessage.Phase, ragMode)
					continue
				}

			} else if chatMessage.Phase == v1pb.Phase_End {
				sendInputMessage(chatMessage.AskId, "", chatMessage.Phase, ragMode)
				continue
			} else if chatMessage.Phase == v1pb.Phase_Close {
				slog.Info("close phase message")
				cancel()
				return
			}

			trimmedContent := strings.TrimSpace(chatMessage.Content)

			if trimmedContent == "" {
				continue
			}

			sendInputMessage(shortuuid.New(), chatMessage.Content, chatMessage.Phase, ragMode)
		}
	}
}

func (s *APIV1Service) handleExchangeMessage(nctx context.Context, cancel context.CancelFunc, tid string, conn *websocket.Conn, outMessageChan chan *v1pb.Exchange) {
	var ackBuilder strings.Builder
	writeMutex := s.WsManager.MutexMap[tid]
	// WebSocket 消息循环
	for {
		select {
		case message, ok := <-outMessageChan:
			if !ok {
				return
			}

			if message.Phase == v1pb.Phase_Continue {
				ackBuilder.WriteString(message.Content)
			}

			if message.Phase == v1pb.Phase_End {
				slog.Info("End phase reached", "UID", message.AskId)

				if ackBuilder.Len() > 0 {
					message.AnswerId = shortuuid.New()
				}

			}

			b, err := proto.Marshal(message)
			if err != nil {
				fmt.Printf("Error marshaling message: %v\n", err)
				return
			}
			fmt.Println("Sending message: ", string(b))
			sendMessage := func() {
				writeMutex.Lock()
				defer writeMutex.Unlock()
				err = conn.WriteMessage(websocket.BinaryMessage, b)

				if err != nil {
					cancel()
					slog.Error("Error writing to WebSocket: %v\n", slog.Any("error", err))
					return
				}
			}
			sendMessage()

		case <-nctx.Done():
			return
		}
	}
}

type AssistantService struct {
	userId int32
	api    *APIV1Service
	ctx    context.Context
	prompt string
	client llm.LLM_Chat
}

func (s *AssistantService) BuildChatMessage(phase v1pb.Phase, content string, askId string) *v1pb.Exchange {
	return &v1pb.Exchange{Phase: phase, Content: content, AskId: askId}
}

func (s *AssistantService) sendOutMessage(ctx context.Context, outMessageChan chan<- *v1pb.Exchange, message *v1pb.Exchange) bool {
	select {
	case <-ctx.Done():
		slog.Warn("chat service context cancelled, stopping message send")
		return false
	default:
		// 如果 context 没有取消，尝试发送消息
		select {
		case outMessageChan <- message:
			return true
		default:
			slog.Warn("channel full or closed")
			return false
		}
	}
}

func (s *AssistantService) doChat(ragMode cpb.RagMode, content string, outMessageChan chan<- *v1pb.Exchange) {
	messages := []llm.RoleContent{}
	messages = append(messages, llm.RoleContent{
		Role:    llm.RoleSys,
		Content: s.prompt,
	})

	messages = append(messages, llm.RoleContent{
		Role:    llm.RoleUser,
		Content: content,
	})

	chatReq := &llm.ChatReq{Msgs: messages}
	chatReq.RspChan = make(chan string)
	chatReq.CompleteChan = make(chan bool)

	go func() {
		select {
		case <-s.ctx.Done():
			slog.Info("recv cancel, return now.")
			return
		default:
			s.sendOutMessage(s.ctx, outMessageChan, s.BuildChatMessage(v1pb.Phase_Begin, ".", ""))
			s.client.Chat(s.ctx, chatReq)
		}
	}()
	var assistantmsgs []string
	for {
		select {
		case <-s.ctx.Done():
			slog.Info("recv cancel. break now.")
			return
		case rsp := <-chatReq.RspChan:
			if s.sendOutMessage(s.ctx, outMessageChan, s.BuildChatMessage(v1pb.Phase_Continue, rsp, "")) {
				break
			}
		case complete := <-chatReq.CompleteChan:
			if complete {
				messages = append(messages, llm.RoleContent{
					Role:    llm.RoleAssistant,
					Content: strings.Join(assistantmsgs, ""),
				})

				assistantmsgs = nil
				if !s.sendOutMessage(s.ctx, outMessageChan, s.BuildChatMessage(v1pb.Phase_End, "", "")) {
					return
				}

				if len(messages) > 110 {
					if !s.sendOutMessage(s.ctx, outMessageChan, s.BuildChatMessage(v1pb.Phase_Close, "记录超过50条，请新建会话", "askid")) {
						return
					}
				}
				return
			}
		case <-time.After(5 * time.Second):
			// 超时处理
		}
	}
}
